package graph

import (
	"traning/algorithm/utility/heap"
)

func PrimMinSum(graph Graph) int {
	var ans int
	nodeSet := map[*Node]bool{}
	edgeSet := map[*Edge]bool{}
	smallHeap := heap.NewSmallHeap(len(graph.Edges))
	for _, node := range graph.Nodes {
		if !nodeSet[node] {
			nodeSet[node] = true
			for _, edge := range node.Edges {
				smallHeap.Push(edge)
				edgeSet[edge] = true
			}
			for smallHeap.Size() > 0 {
				edge := smallHeap.Pop().(*Edge)
				if !nodeSet[edge.To] {
					ans += edge.Weight
					nodeSet[edge.To] = true
					for _, nextEdge := range edge.To.Edges {
						if !edgeSet[nextEdge] {
							smallHeap.Push(nextEdge)
						}
					}
				}
			}
		}
	}
	return ans
}